import os
import numpy as np
from PIL import Image
from tqdm.auto import tqdm

def create_npz_from_cifar10(sample_dir, output_path, max_images_per_class=None):
    """
    Builds a single .npz file from a folder of .png samples organized by class.
    :param sample_dir: Directory containing class folders (e.g., 'data/cifar10/raw')
    :param output_path: Path to save the .npz file
    :param max_images_per_class: Optional. Limit the number of images per class.
    """
    samples = []
    class_folders = sorted([d for d in os.listdir(sample_dir) if os.path.isdir(os.path.join(sample_dir, d))])

    for class_name in class_folders:
        class_dir = os.path.join(sample_dir, class_name)
        image_files = sorted([f for f in os.listdir(class_dir) if f.endswith('.png')])
        
        if max_images_per_class is not None:
            image_files = image_files[:max_images_per_class]

        for image_file in tqdm(image_files, desc=f"Processing class {class_name}"):
            image_path = os.path.join(class_dir, image_file)
            sample_pil = Image.open(image_path)
            sample_np = np.asarray(sample_pil).astype(np.uint8)
            samples.append(sample_np)

    samples = np.stack(samples)
    print(f"Collected {len(samples)} images with shape {samples.shape[1:]}")

    np.savez(output_path, arr_0=samples)
    print(f"Saved .npz file to {output_path} [shape={samples.shape}].")
    return output_path

# Example usage
sample_dir = "/home/zuwenqiang/Respo2/SRA/preprocessing/data/cifar10/raw"
output_path = "samples/ref/cifar10.npz"
create_npz_from_cifar10(sample_dir, output_path)